from PyQt6.QtCore import QByteArray, QBuffer, QIODevice
from PyQt6.QtGui import QImage
import numpy as np

def numpy_to_qimage(numpy_image):
    """
    将 numpy 数组转换为 QImage 对象
    :param numpy_image: numpy 数组，形状为 (height, width, channels)
    :return: QImage 对象
    """

    if numpy_image is None:
        print("numpy 数组为空，无法转换为 QImage")
        return None

    # 检查是否为 numpy 数组
    if not isinstance(numpy_image, np.ndarray):
        print("不是 numpy 数组，无法转换为 QImage")
        return None

    if len(numpy_image.shape) == 2:
        # 灰度图像
        channels = 1
        height, width = numpy_image.shape
    elif len(numpy_image.shape) == 3:
        # RGB 图像或 RGBA 图像
        height, width, channels = numpy_image.shape
    else:
        print("不支持的图像形状：必须是 2D（灰度）或 3D（RGB 或 RGBA）")
        return None

    # 根据通道数确定 QImage 的格式
    if channels == 1:
        # 灰度图像
        qimage_format = QImage.Format.Format_Grayscale8
    elif channels == 3:
        # RGB 图像
        qimage_format = QImage.Format.Format_RGB888
    elif channels == 4:
        # RGBA 图像
        qimage_format = QImage.Format.Format_RGBA8888
    else:
        raise ValueError("不支持的图像通道数：必须是 1（灰度）、3（RGB）或 4（RGBA）")

    trans_qimage = QImage(numpy_image.data, width, height, numpy_image.strides[0], qimage_format)

    # 将 numpy 数组转换为 QImage
    return trans_qimage

def qimage_to_bytes(input_qimage, image_format='PNG'):
    """将 QImage 转换为字节流"""
    if input_qimage.isNull():
        print("QImage 无效")
        return None

    byte_array = QByteArray()
    buffer = QBuffer(byte_array)
    # 打开 QBuffer
    if not buffer.open(QIODevice.OpenModeFlag.ReadWrite):
        print("QBuffer 打开失败")
        return None

    # 保存 QImage 到 QBuffer
    if not input_qimage.save(buffer, image_format):
        print("保存 QImage 失败")
        buffer.close()
        return None
    buffer.close()
    return byte_array.data()


# 示例用法
if __name__ == "__main__":
    # 假设 numpy_images 是一个包含图像数据的 numpy 数组
    numpy_images = np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8)  # 生成一个随机 RGB 图像

    # 转换为 QImage
    qimage = numpy_to_qimage(numpy_images)

    # 检查 QImage 是否有效
    if qimage.isNull():
        print("转换失败：QImage 无效")
    else:
        print("转换成功：QImage 大小:", qimage.size())
